#!/bin/bash
# Authored : <fanyuantao>
# Description: This script is used to setup ssh key for remote access.

server=""

function ssh_keygen_check() {
    if [ -f ~/.ssh/id_rsa ]; then
        echo "SSH key already exists. Do you want to recreate it? (y/n) "
        read yn
        if [ "$yn" == "y" ]; then
            ssh-keygen -t rsa -b 4096 -C "fanyuantao@greatech.com.cn" -f ~/.ssh/id_rsa
        fi
    else
        echo "Generating SSH key..."
         # generate ssh key
        ssh-keygen -t rsa -b 4096 -C "fanyuantao@greatech.com.cn" -f ~/.ssh/id_rsa
    fi
}

function ssh_config() {
    while true; do
        echo ""
        echo "1. Check SSH key"
        echo "2. Copy SSH key to remote server"
        echo "3. Check remote server ssh access"
        echo "0. Exit"
        read -p "Please select an option : " selection
        case $selection in
            1)
                ssh_keygen_check
                ;;
            2)
                ssh_copy_id
                ;;
            3)
                remote_server_ssh_check
                ;;
            0)
                exit 0
                ;;
            *)
                echo "Invalid option. Please try again."
                ;;
        esac
    done
}

function remote_server_ssh_check() {
    input_server_info
    if [[ -z $server || ! $(ssh -o ConnectTimeout=5 $server "echo success") ]]; then
        echo "Remote server $server is not accessible."
        return 1
    else
        echo "Remote server $server is accessible."
    fi
}

function input_server_info() {
    echo ""
    echo "example: 192.168.1.100"
    read -p "Input remote server IP address: " ip
    if [[ -z $ip || ! $ip =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
        echo "Invalid input. Please enter the correct remote address."
        return 1
    else
        server_ip=$ip
    fi

    echo ""
    read -p "Input remote server username of $server_ip : " username
    if [[ -z $username ]]; then
        echo "Invalid input. Please enter the correct remote username."
        return 1
    else
        server_username=$username
    fi

    server="$server_username@$server_ip"

    echo $server
}

function check_ping_command() {
    if ! command -V ping >/dev/null 2>&1; then
        echo "ping command is not installed. Please install it first."
        exit 1
    fi
    # CentOS/AnolisOS 安装 ping 命令
    # yum install -y -q iputils
    # Ubuntu 安装 ping 命令
    # sudo apt-get install -y -q iputils-ping
}

function ssh_copy_id() {
    input_server_info
    if ping -c 4 $server_ip >/dev/null 2>&1; then
        echo "Remote server $server_ip is accessible."
        echo "Starting copy ssh key to remote : $server "
        ssh-copy-id -i ~/.ssh/id_rsa.pub $server
    else
        echo "Remote server $server_ip is not accessible. Please check the IP address or username."
        return 1
    fi
}

function main() {
    check_ping_command
    ssh_config
}

main